Add forward LSE output to FlyDSL flash attention kernels#844
Add forward LSE output to FlyDSL flash attention kernels#844amd-wsung102 wants to merge 1 commit into
Conversation
|
Thanks for adding the forward LSE support. A couple of requests before this can be merged:
|
| exe(q_flat, k_flat, v_flat, o_flat, B, Sq, **kwargs) | ||
| else: | ||
| kwargs: dict = dict(stream=launch_stream) | ||
| kwargs: dict = dict(lse=lse, stream=launch_stream) |
There was a problem hiding this comment.
Blocker: this passes lse to every dense launcher, including the fp8 path when return_lse=False. The fp8 launcher does not accept an lse keyword, so fp8 forward will fail before launch with an unexpected keyword argument. Please only add lse to kwargs for non-fp8 launchers, or update the fp8 launcher to accept and ignore lse=None.
| cq = cu_seqlens_q if cu_seqlens_q is not None else Out | ||
| ck = cu_seqlens_kv if cu_seqlens_kv is not None else Out | ||
| bt = block_table if block_table is not None else Out | ||
| ls = lse if lse is not None else Out |
There was a problem hiding this comment.
Blocker: when this builder is created with return_lse=True but the caller omits lse, this falls back to Out. Since RETURN_LSE stores fp32 LSE, that can overwrite the output buffer or write through an incompatible layout. Please raise ValueError when RETURN_LSE and lse is None instead of using Out as the fallback.
| block_table_stride = 0 | ||
| # LSE is only written under const_expr(RETURN_LSE); O placeholder otherwise. | ||
| if lse is None: | ||
| lse = O |
There was a problem hiding this comment.
Blocker: same issue here. If return_lse=True and the caller forgets to pass lse, this falls back to O, so the kernel can write fp32 LSE into the output buffer. Please fail fast with ValueError when RETURN_LSE and lse is None.
|
One style request: for code comments, except function-level comments/docstrings, please keep comments concise, ideally 1-2 lines. Several added inline comments are longer than needed and would be easier to maintain if shortened. |
2 participants
Motivation
This PR is motivated by the task described here https://amd-hub.atlassian.net/jira/software/c/projects/AIMODELS/boards/2992?assignee=712020%3Acedccbdf-bb1e-4e04-8c3b-edcdd7744eda&selectedIssue=AIMODELS-983, which prepares for a MSLK integration.
Summary
Extends the FlyDSL flash attention forward kernels to optionally emit the per-row log-sum-exp (LSE) alongside the attention output. LSE is the normalization statistic the backward pass needs to recompute the softmax probabilities
Pwithout re-running the online softmax, so this is the enablement step for a pure-FlyDSL forward+backward attention path.Opt-in via
return_lse=True; default behavior is unchanged.Changes
flash_attn_interface.py: addreturn_lsetoflydsl_flash_attn_func; thread it through every build-cache helper; allocate fp32[B, num_heads, Sq]and return(out, lse)when enabled. Cache is keyed onreturn_lseso the existing no-LSE kernels are never perturbed.flash_attn_generic.py:return_lsebuild flag;LSEtensor in the kernel/launch/compile signatures; storesm_scale*m_raw + ln(l)on both the normal store and the zeroed (fully-masked) tile.flash_attn_gfx950.py: same for the dualwave kernel (dense storem_row*ln2 + ln(l_row)), plus the split-K combine kernel finalizingm_max*ln2 + ln(den)from the per-split(m, l). One writer per row via half-wave + OOB-sentinel guards.tests/kernels/test_flash_attn_lse.py: numerical validation vs a float32 PyTorch reference.docs/flash_attn_lse_contract.md: the A1↔A3 interface contract.Output tensor
float32, shape[B, num_heads, Sq](varlen:[B, num_heads, max_seqlen_q], padded rows undefined).Supported / unsupported paths
num_kv_splits > 1)cu_seqlens)NotImplementedErrorblock_table/seqlen_k)NotImplementedErrorfp8 and paged KV are decode-oriented and out of scope for the training backward path; they fail loudly rather than silently returning wrong LSE.
Validation
tests/kernels/test_flash_attn_lse.py— 33 cases pass on gfx950, covering dense MHA/GQA/MQA (D=64/128, bf16/f16), split-K, varlen, cross-attention, and fully-masked (-inf) rows. LSE matches the fp32 reference within ~8e-3 (bf16) / ~4e-3 (f16); fp32-accumulating paths (varlen light,D=64) match to ~1e-6.Performance — LSE store overhead
The LSE store is one fp32 write per query row (single lane per row), negligible next to the
O(S²·D)attention compute. Measured on gfx950 (MI350), bf16, H=16, D=128, min-of-trials:torch.emptyallocation of the LSE output tensor, not the GPU store.Performance — FlyDSL forward vs CK
FlyDSL forward throughput vs the CK backend (aiter
mha_fwd), confirming the pure-FlyDSL forward path is competitive. gfx950 (MI350), bf16, host-side end-to-end (CUDA events), mean ± stddev over 3 runs.aiter_asm(hand-tunedfmha_v3) shown for reference. Baseline FlyDSL forward (LSE adds ≤1.3%, see above).FlyDSL is faster than CK on every shape measured (1.17×–1.59×), and matches or beats the hand-tuned
aiter_asmpath on non-causal, GQA, short-sequence, andD=64shapes (asm remains ahead on long causal). Correctness matches CK: FlyDSL max-err vs the shared reference is on par (Fly/CK max-err ratio = 1.00×; e.g. 3.9e-03 bf16 causal, 2.4e-04 non-causal), all runs PASS the harness gate.Compatibility
Fully backwards compatible —
return_lsedefaults toFalseand returns a bare output tensor; LSE build variants are cached separately.Test plan
pytest tests/kernels/test_flash_attn_lse.pyon gfx950 (33 passed)